import pandas as pd
from ultralytics import YOLO
from mlflow.pyfunc import PythonModel

class YOLOModel(PythonModel):
    def load_context(self, context):
        self.model = YOLO("path/to/weights.pt")

    def predict(self, context, inputs):
        input_data = inputs.values
        results = self.model(input_data)
        outputs = results.xyxy[0].numpy()
        dic = {"class": outputs[:, 5], "confidence": outputs[:, 4], "x1": outputs[:, 0], "y1": outputs[:, 1], "x2": outputs[:, 2], "y2": outputs[:, 3]}
        predict_value = pd.DataFrame(dic, index=range(0, len(outputs)))
        return predict_value